package org.finesys.chat.core.base.embedding.store.milvus;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

import org.finesys.chat.core.base.embedding.filter.Filter;
import org.finesys.chat.core.base.embedding.store.milvus.filter.MilvusMetadataFilterMapper;

import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.param.MetricType;
import io.milvus.param.collection.DropCollectionParam;
import io.milvus.param.collection.FlushParam;
import io.milvus.param.collection.HasCollectionParam;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.QueryParam;
import io.milvus.param.dml.SearchParam;

public class CollectionRequestBuilder {
    CollectionRequestBuilder() {
    }

    static FlushParam buildFlushRequest(String collectionName) {
        return FlushParam.newBuilder().withCollectionNames(Collections.singletonList(collectionName)).build();
    }

    static HasCollectionParam buildHasCollectionRequest(String collectionName) {
        return HasCollectionParam.newBuilder().withCollectionName(collectionName).build();
    }

    static DropCollectionParam buildDropCollectionRequest(String collectionName) {
        return DropCollectionParam.newBuilder().withCollectionName(collectionName).build();
    }

    static InsertParam buildInsertRequest(String collectionName, List<InsertParam.Field> fields) {
        return InsertParam.newBuilder().withCollectionName(collectionName).withFields(fields).build();
    }

    static LoadCollectionParam buildLoadCollectionInMemoryRequest(String collectionName) {
        return LoadCollectionParam.newBuilder().withCollectionName(collectionName).build();
    }

    static SearchParam buildSearchRequest(String collectionName, FieldDefinition fieldDefinition, List<Double> vector, Filter filter, int maxResults, MetricType metricType, ConsistencyLevelEnum consistencyLevel) {
        SearchParam.Builder builder = SearchParam.newBuilder().withCollectionName(collectionName).withVectors(Collections.singletonList(vector)).withVectorFieldName(fieldDefinition.getVectorFieldName()).withTopK(maxResults).withMetricType(metricType).withConsistencyLevel(consistencyLevel).withOutFields(Arrays.asList(fieldDefinition.getIdFieldName(), fieldDefinition.getTextFieldName(), fieldDefinition.getMetadataFieldName()));
        if (filter != null) {
            builder.withExpr(MilvusMetadataFilterMapper.map(filter, fieldDefinition.getMetadataFieldName()));
        }

        return builder.build();
    }

    static QueryParam buildQueryRequest(String collectionName, FieldDefinition fieldDefinition, List<String> rowIds, ConsistencyLevelEnum consistencyLevel) {
        return QueryParam.newBuilder().withCollectionName(collectionName).withExpr(buildQueryExpression(rowIds, fieldDefinition.getIdFieldName())).withConsistencyLevel(consistencyLevel).withOutFields(Collections.singletonList(fieldDefinition.getVectorFieldName())).build();
    }

    static DeleteParam buildDeleteRequest(String collectionName, String expr) {
        return DeleteParam.newBuilder().withCollectionName(collectionName).withExpr(expr).build();
    }

    private static String buildQueryExpression(List<String> rowIds, String idFieldName) {
        return (String) rowIds.stream().map((id) -> {
            return String.format("%s == '%s'", idFieldName, id);
        }).collect(Collectors.joining(" || "));
    }
}
